#pragma once

#include "IAESEngine.h"
#include "AESkey.h"
#include "AESUserDefs.h"
#include "AESMath.h"

class AESEngine : public IAESEngine
{
public:
	AESEngine(void);
	virtual ~AESEngine(void);

// APIs supporting the published interface, IAESEngine 
public:
	virtual int Encrypt(IAESKey* key, const byte in[], byte out[]);
	virtual int Decrypt(IAESKey* key, const byte in[], byte out[]);

private:
	// validates state 
	bool	validateState		(void);
	int		ConfigureState		(IAESKey* key);
	void    CloseState			(void);
	void	setup				(void);

	// building blocks for our cipher
	int		ExpandKey			(void);
	void	SubBytes			(byte** state);
	void	ShiftRows			(byte** state);
	void	MixColumns			(byte** s);
	void	AddRoundKey			(byte** state);
	void	Encipher			(byte in[], byte out[]); 

	// Helpers
	bool	verifyKeyIdentity	(IAESKey* _key);
	void	setKeyLength		(IAESKey* key);
	void	setNumberRounds		(IAESKey* key);
	void	buildMathTables		(void);
	int		buildKeyTable		(void);
	int		buildStateTable		(void);
	void	removeMathTables	(void);
	void	removeKeyTable		(void);
	void	removeStateTable	(void);
	void	Copy				(byte* out, byte** state); 
	void	Copy				(byte** state, byte* in); 
	int		CopyKey				(void);
	int		newAESWordBuffer	(byte **outbuffer);


private:
	// helpers to print
	byte	Hex					(byte a);
	void	printArray			(char name[], byte a[], int length);
	void	printArray			(char name[], int round, byte** s);




private:

	// words in a block, always 4 for now
	static const int	Nb = AES_NUM_WORDS_IN_BLOCK; 

	int					Nk;				// key length in words
	int					Nr;				// number of rounds, = Nk + 6
	int					_wCount;		// position in w for RoundKey (= 0 each encrypt)
	byte*				_w;				// the expanded key
	AESMath*			_mathTables;	// all the tables needed for AES
	byte**				_state;
	IAESKey*			_key;
};

